#include <winsock2.h>     
#include <iostream>       
#include <thread>         
#include <vector>         
#include <ws2tcpip.h>    
#include <getopt.h>       // 用于命令行参数解析
#pragma comment(lib,"ws2_32.lib") 

using namespace std;

// 获取当前时间字符串，用于日志记录
string getCurrentTime()
{
    std::time_t now = std::time(nullptr);    
    std::tm* localTime = std::localtime(&now); 

    char timeStr[80];
    strftime(timeStr, sizeof(timeStr), "%Y-%m-%d %H:%M:%S", localTime);
    return std::string(timeStr);
}

// 解析SOCK5请求，获取目标地址和端口
bool parseSOCK5Request(SOCKET clientSock, string& targetIP, int& targetPort, bool& isDomain) {
    char buffer[1024];
    int recvLen = recv(clientSock, buffer, sizeof(buffer), 0);
    
    if (recvLen < 7) {
        cerr << getCurrentTime() << " SOCK5请求格式错误，长度不足" << endl;
        return false;
    }
    
    // 检查协议版本
    if (buffer[0] != 5) {
        cerr << getCurrentTime() << " 非SOCK5协议，版本号: " << (int)buffer[0] << endl;
        return false;
    }
    
    // 响应SOCK5握手（选择无认证方式）
    char handshakeResp[2] = {5, 0}; // 版本号+无认证
    if (send(clientSock, handshakeResp, 2, 0) == SOCKET_ERROR) {
        cerr << getCurrentTime() << " SOCK5握手响应失败" << endl;
        return false;
    }
    
    // 再次接收请求（实际请求数据）
    recvLen = recv(clientSock, buffer, sizeof(buffer), 0);
    if (recvLen < 7) {
        cerr << getCurrentTime() << " SOCK5实际请求格式错误" << endl;
        return false;
    }
    
    // 解析目标地址类型
    uint8_t addrType = buffer[3];
    isDomain = (addrType == 3); // 3表示域名地址
    
    // 解析目标端口（网络字节序转主机字节序）
    targetPort = ntohs(*(uint16_t*)(buffer + 1));
    
    // 解析目标IP或域名
    if (addrType == 1) { // IPv4地址
        struct in_addr addr;
        memcpy(&addr, buffer + 4, 4);
        targetIP = inet_ntoa(addr);
    } else if (addrType == 3) { // 域名
        int domainLen = buffer[4];
        if (recvLen < 5 + domainLen) {
            cerr << getCurrentTime() << " 域名长度错误" << endl;
            return false;
        }
        buffer[5 + domainLen] = '\0';
        targetIP = string((char*)(buffer + 5), domainLen);
    } else if (addrType == 4) { // IPv6地址（暂不支持）
        cerr << getCurrentTime() << " 暂不支持IPv6地址" << endl;
        return false;
    } else {
        cerr << getCurrentTime() << " 未知地址类型: " << (int)addrType << endl;
        return false;
    }
    
    cout << getCurrentTime() << " 解析SOCK5请求: " 
         << (isDomain ? "域名: " : "IP: ") << targetIP 
         << ", 端口: " << targetPort << endl;
    return true;
}

// 解析SOCK4请求（简化实现）
bool parseSOCK4Request(SOCKET clientSock, string& targetIP, int& targetPort, bool& isDomain) {
    char buffer[1024];
    int recvLen = recv(clientSock, buffer, sizeof(buffer), 0);
    
    if (recvLen < 8) {
        cerr << getCurrentTime() << " SOCK4请求格式错误，长度不足" << endl;
        return false;
    }
    
    // 检查协议版本
    if (buffer[0] != 4) {
        cerr << getCurrentTime() << " 非SOCK4协议，版本号: " << (int)buffer[0] << endl;
        return false;
    }
    
    // 解析目标端口（网络字节序转主机字节序）
    targetPort = ntohs(*(uint16_t*)(buffer + 2));
    
    // 解析目标IP
    struct in_addr addr;
    memcpy(&addr, buffer + 4, 4);
    targetIP = inet_ntoa(addr);
    isDomain = false;
    
    // 检查是否有用户名（SOCK4a支持域名）
    if (targetIP == "0.0.0.0") {
        // SOCK4a格式，域名在用户名位置
        size_t usernamePos = 8;
        while (usernamePos < recvLen && buffer[usernamePos] != 0) {
            usernamePos++;
        }
        if (usernamePos < recvLen - 1) {
            targetIP = string((char*)(buffer + 8), usernamePos - 8);
            isDomain = true;
        }
    }
    
    cout << getCurrentTime() << " 解析SOCK4请求: " 
         << (isDomain ? "域名: " : "IP: ") << targetIP 
         << ", 端口: " << targetPort << endl;
    return true;
}

// 域名解析函数
bool resolveDomain(const string& domain, string& ipAddr) {
    struct addrinfo hints, *result, *ptr;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET;       // 只获取IPv4地址
    hints.ai_socktype = SOCK_STREAM;
    
    int status = getaddrinfo(domain.c_str(), NULL, &hints, &result);
    if (status != 0) {
        cerr << getCurrentTime() << " 域名解析失败: " << gai_strerror(status) << endl;
        return false;
    }
    
    // 遍历结果，获取第一个IPv4地址
    bool found = false;
    for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
        if (ptr->ai_family == AF_INET) {
            struct sockaddr_in* ipv4 = (struct sockaddr_in*)ptr->ai_addr;
            ipAddr = inet_ntoa(ipv4->sin_addr);
            found = true;
            break;
        }
    }
    
    freeaddrinfo(result); // 释放资源
    return found;
}

// 客户端处理线程函数
void handleClient(SOCKET socketClient)
{
    sockaddr_in addrClient; 
    int lenAddr = sizeof(addrClient);
    getpeername(socketClient, (struct sockaddr*)&addrClient, &lenAddr);

    char ipClient[INET_ADDRSTRLEN]; 
    inet_ntop(AF_INET, &(addrClient.sin_addr), ipClient, INET_ADDRSTRLEN);

    cout << getCurrentTime() << " 客户端 " << ipClient << " 连接成功" << endl;

    string targetIP;
    int targetPort = 0;
    bool isDomain = false;
    bool isSOCK5 = false;
    
    // 读取第一个字节判断协议版本
    char versionBuf[1];
    int recvLen = recv(socketClient, versionBuf, 1, 0);
    if (recvLen <= 0) {
        cerr << getCurrentTime() << " 客户端 " << ipClient << " 接收协议版本失败" << endl;
        closesocket(socketClient);
        return;
    }
    
    // 判断协议类型并解析请求
    if (versionBuf[0] == 5) {
        isSOCK5 = true;
        if (!parseSOCK5Request(socketClient, targetIP, targetPort, isDomain)) {
            closesocket(socketClient);
            return;
        }
    } else if (versionBuf[0] == 4) {
        if (!parseSOCK4Request(socketClient, targetIP, targetPort, isDomain)) {
            closesocket(socketClient);
            return;
        }
    } else {
        cerr << getCurrentTime() << " 未知协议版本: " << (int)versionBuf[0] << endl;
        closesocket(socketClient);
        return;
    }
    
    // 处理域名解析
    if (isDomain) {
        string resolvedIP;
        if (!resolveDomain(targetIP, resolvedIP)) {
            cerr << getCurrentTime() << " 域名 " << targetIP << " 解析失败" << endl;
            // 向客户端发送SOCK5连接失败响应
            if (isSOCK5) {
                char failureResp[10] = {5, 4, 0, 1, 0, 0, 0, 0, 0, 0}; // 版本+失败状态+IPv4地址
                send(socketClient, failureResp, 10, 0);
            }
            closesocket(socketClient);
            return;
        }
        targetIP = resolvedIP;
    }
    
    cout << getCurrentTime() << " 客户端 " << ipClient << " 请求连接至：" << targetIP << ":" << targetPort << endl;

    // 创建连接目标服务器的套接字
    SOCKET socketTarget = socket(AF_INET, SOCK_STREAM, 0); 
    if (socketTarget == INVALID_SOCKET) {
        cerr << getCurrentTime() << " 为客户端 " << ipClient << " 创建目标套接字失败: " << WSAGetLastError() << endl;
        closesocket(socketClient);
        return;
    }

    sockaddr_in addrTarget = {}; 
    addrTarget.sin_family = AF_INET; 
    addrTarget.sin_port = htons(targetPort); 
    
    // 转换目标IP地址
    if (inet_pton(AF_INET, targetIP.c_str(), &addrTarget.sin_addr) <= 0) {
        cerr << getCurrentTime() << " 无效的目标IP地址: " << targetIP << endl;
        closesocket(socketClient);
        closesocket(socketTarget);
        return;
    }

    // 连接到目标服务器
    if (connect(socketTarget, (struct sockaddr*)&addrTarget, sizeof(addrTarget)) == SOCKET_ERROR) {
        cerr << getCurrentTime() << " 客户端 " << ipClient << " 连接目标 " << targetIP << ":" << targetPort << " 失败: " << WSAGetLastError() << endl;
        // 发送SOCK5连接失败响应
        if (isSOCK5) {
            char failureResp[10] = {5, 4, 0, 1, 0, 0, 0, 0, 0, 0}; // 版本+失败状态+IPv4地址
            send(socketClient, failureResp, 10, 0);
        }
        closesocket(socketClient);
        closesocket(socketTarget);
        return;
    }

    cout << getCurrentTime() << " 客户端 " << ipClient << " 成功连接到目标 " << targetIP << ":" << targetPort << endl;
    
    // 发送SOCK5连接成功响应
    if (isSOCK5) {
        char successResp[10] = {5, 0, 0, 1, 0, 0, 0, 0, 0, 0}; // 版本+成功状态+IPv4地址
        if (send(socketClient, successResp, 10, 0) == SOCKET_ERROR) {
            cerr << getCurrentTime() << " 向客户端 " << ipClient << " 发送成功响应失败: " << WSAGetLastError() << endl;
            closesocket(socketClient);
            closesocket(socketTarget);
            return;
        }
    }

    // 设置用于select函数的文件描述符集合
    fd_set fdsRead;
    vector<SOCKET> listSockets; 
    listSockets.push_back(socketClient);
    listSockets.push_back(socketTarget);

    char bufTransfer[4096]; 
    int bytesRead, bytesSent;

    // 进入数据双向转发循环
    while (true) {
        FD_ZERO(&fdsRead); 
        for (size_t i = 0; i < listSockets.size(); ++i) {
            FD_SET(listSockets[i], &fdsRead);
        }

        if (listSockets.empty()) {
            cout << getCurrentTime() << " 客户端 " << ipClient << " 与目标 " << targetIP << ":" << targetPort << " 的连接均已关闭" << endl;
            break;
        }

        int activity = select(0, &fdsRead, NULL, NULL, NULL);
        if (activity == SOCKET_ERROR) {
            cerr << getCurrentTime() << " select 函数出错: " << WSAGetLastError() << endl;
            break;
        }

        for (size_t i = 0; i < listSockets.size(); ) { 
            SOCKET s = listSockets[i];
            if (FD_ISSET(s, &fdsRead)) { 
                bytesRead = recv(s, bufTransfer, sizeof(bufTransfer), 0); 
                if (bytesRead <= 0) { 
                    if (bytesRead == 0) {
                        cout << getCurrentTime() << " 套接字 " << (s == socketClient ? ipClient : targetIP) << " 连接已关闭" << endl;
                    } else {
                        cerr << getCurrentTime() << " 从套接字 " << (s == socketClient ? ipClient : targetIP) << " 接收数据失败: " << WSAGetLastError() << endl;
                    }
                    closesocket(s); 
                    listSockets.erase(listSockets.begin() + i);
                    SOCKET otherSocket = (s == socketClient) ? socketTarget : socketClient;
                    auto it = find(listSockets.begin(), listSockets.end(), otherSocket);
                    if (it != listSockets.end()) {
                        cout << getCurrentTime() << " 关闭配对套接字 " << (otherSocket == socketClient ? ipClient : targetIP) << endl;
                        closesocket(otherSocket);
                        listSockets.erase(it);
                    }
                    goto end_inner_loop;
                } else { 
                    SOCKET socketSendTo = (s == socketClient) ? socketTarget : socketClient;
                    bool canSend = false;
                    for(const auto& activeSock : listSockets) {
                        if (activeSock == socketSendTo) {
                            canSend = true;
                            break;
                        }
                    }

                    if (canSend) {
                        bytesSent = send(socketSendTo, bufTransfer, bytesRead, 0); 
                        if (bytesSent <= 0) {
                            cerr << getCurrentTime() << " 发送数据到 " << (socketSendTo == socketClient ? ipClient : targetIP) << " 失败: " << WSAGetLastError() << endl;
                            goto cleanup_and_exit_thread;
                        }
                    } else {
                        cout << getCurrentTime() << " 目标发送套接字已关闭，数据未发送" << endl;
                    }
                    i++;
                }
            } else {
                i++;
            }
        }
        end_inner_loop:;
    }

cleanup_and_exit_thread:
    for (SOCKET s_cleanup : listSockets) {
        closesocket(s_cleanup);
    }
    listSockets.clear();

    cout << getCurrentTime() << " 客户端 " << ipClient << " 处理线程结束" << endl;
}

int main(int argc, char* argv[])
{
    WSADATA dataWSA; 
    SOCKET socketServer, socketClient; 
    sockaddr_in addrServer, addrClient; 
    int lenClientAddr = sizeof(addrClient);
    int port = 1080; // 默认端口设为1080

    // 解析命令行参数
    int opt;
    while ((opt = getopt(argc, argv, "n:")) != -1) {
        switch (opt) {
            case 'n':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    cerr << "端口号必须在1-65535之间" << endl;
                    return -1;
                }
                break;
            default:
                cerr << "用法: " << argv[0] << " -n <端口号>" << endl;
                return -1;
        }
    }

    // 1. 初始化Winsock
    if (WSAStartup(MAKEWORD(2, 2), &dataWSA) != 0) 
    {
        cerr << "Winsock 初始化失败，错误代码: " << WSAGetLastError() << endl;
        return -1;
    }

    // 2. 创建服务器监听套接字
    socketServer = socket(AF_INET, SOCK_STREAM, 0); 
    if (socketServer == INVALID_SOCKET) 
    {
        cerr << "套接字创建失败，错误代码: " << WSAGetLastError() << endl;
        WSACleanup(); 
        return -1;
    }

    // 3. 配置服务器地址结构（监听所有接口）
    addrServer.sin_family = AF_INET; 
    addrServer.sin_addr.s_addr = INADDR_ANY; // 监听所有网络接口
    addrServer.sin_port = htons(port); 

    // 4. 绑定套接字到服务器地址和端口
    if (bind(socketServer, (struct sockaddr*)&addrServer, sizeof(addrServer)) == SOCKET_ERROR)
    {
        cerr << "套接字绑定失败，错误代码: " << WSAGetLastError() << endl;
        closesocket(socketServer); 
        WSACleanup(); 
        return -1;
    }

    // 5. 开始监听连接请求
    if (listen(socketServer, SOMAXCONN) == SOCKET_ERROR)
    {
        cerr << "监听失败，错误代码: " << WSAGetLastError() << endl;
        closesocket(socketServer);
        WSACleanup();
        return -1;
    }

    // 转换端口为字符串用于日志
    char portStr[6];
    sprintf(portStr, "%d", port);
    cout << getCurrentTime() << " 服务端在 0.0.0.0:" << portStr << " 等待连接..." << endl;

    // 6. 循环接受客户端连接
    while (true)
    {
        // accept会阻塞，直到有客户端连接请求到达
        socketClient = accept(socketServer, (struct sockaddr*)&addrClient, &lenClientAddr);
        if (socketClient == INVALID_SOCKET) 
        {
            cerr << "连接接受失败，错误代码: " << WSAGetLastError() << endl;
            continue; // 继续等待下一个连接
        }

        // 获取并打印新连接的客户端信息
        char clientIpStr[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &addrClient.sin_addr, clientIpStr, INET_ADDRSTRLEN);
        cout << getCurrentTime() << " 接受到来自 " << clientIpStr << ":" << ntohs(addrClient.sin_port) << " 的连接" << endl;

        // 7. 为每个客户端创建一个新线程进行处理
        try {
            thread threadClient(handleClient, socketClient); 
            threadClient.detach(); 
        } catch (const exception& e) {
            cerr << "创建线程失败: " << e.what() << endl;
            closesocket(socketClient);
        }
    }

    // 清理资源（正常情况下不会执行到这里）
    closesocket(socketServer);
    WSACleanup();

    return 0;
}